import torch
import torch.nn as nn
import time
from .utils import calculate_accuracy, save_model, count_parameters
from .tensorboard_logger import TensorBoardLogger


class Trainer:
    """通用深度学习训练器"""
    
    def __init__(self, model, device, criterion=None, optimizer=None, scheduler=None,
                 tensorboard_logger=None, model_save_path=None, print_freq=50):
        """
        初始化训练器
        
        Args:
            model (torch.nn.Module): 模型
            device (torch.device): 计算设备
            criterion (torch.nn.Module): 损失函数
            optimizer (torch.optim.Optimizer): 优化器
            scheduler (torch.optim.lr_scheduler): 学习率调度器
            tensorboard_logger (TensorBoardLogger): TensorBoard记录器
            model_save_path (str): 模型保存路径
            print_freq (int): 打印频率
        """
        self.model = model.to(device)
        self.device = device
        self.criterion = criterion if criterion else nn.CrossEntropyLoss()
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.tensorboard_logger = tensorboard_logger
        self.model_save_path = model_save_path
        self.print_freq = print_freq
        
        # 统计模型参数
        param_count = count_parameters(self.model)
        print(f'\n模型参数数量: {param_count:,}')
        
        # 记录模型结构到TensorBoard
        if self.tensorboard_logger:
            try:
                # 假设输入为224x224的3通道图像
                dummy_input = torch.randn(1, 3, 224, 224).to(device)
                self.tensorboard_logger.log_model_graph(self.model, dummy_input)
            except Exception as e:
                print(f'记录模型结构时出错: {e}')
    
    def train_epoch(self, train_loader, epoch):
        """
        训练一个epoch
        
        Args:
            train_loader (DataLoader): 训练数据加载器
            epoch (int): 当前epoch
        
        Returns:
            tuple: (平均损失, 平均准确率)
        """
        self.model.train()
        running_loss = 0.0
        running_corrects = 0
        total_samples = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            # 前向传播
            self.optimizer.zero_grad()
            outputs = self.model(data)
            loss = self.criterion(outputs, target)
            
            # 反向传播
            loss.backward()
            self.optimizer.step()
            
            # 统计
            running_loss += loss.item()
            batch_acc = calculate_accuracy(outputs, target)
            running_corrects += batch_acc * target.size(0)
            total_samples += target.size(0)
            
            # 记录到TensorBoard
            if self.tensorboard_logger:
                global_step = epoch * len(train_loader) + batch_idx
                self.tensorboard_logger.log_training_batch(global_step, loss.item(), batch_acc * 100)
            
            # 打印进度
            if batch_idx % self.print_freq == 0:
                current_acc = 100. * running_corrects / total_samples
                print(f'  批次 [{batch_idx}/{len(train_loader)}], '
                      f'损失: {loss.item():.4f}, '
                      f'准确率: {current_acc:.2f}%')
        
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * running_corrects / total_samples
        return epoch_loss, epoch_acc
    
    def validate(self, val_loader):
        """
        验证模型
        
        Args:
            val_loader (DataLoader): 验证数据加载器
        
        Returns:
            tuple: (平均损失, 平均准确率)
        """
        self.model.eval()
        val_loss = 0.0
        running_corrects = 0
        total_samples = 0
        
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(self.device), target.to(self.device)
                outputs = self.model(data)
                loss = self.criterion(outputs, target)
                
                val_loss += loss.item()
                batch_acc = calculate_accuracy(outputs, target)
                running_corrects += batch_acc * target.size(0)
                total_samples += target.size(0)
        
        val_loss /= len(val_loader)
        val_acc = 100. * running_corrects / total_samples
        return val_loss, val_acc
    
    def test(self, test_loader, train_dataset=None, label_mapping=None, num_classes=100):
        """
        测试模型
        
        Args:
            test_loader (DataLoader): 测试数据加载器
            train_dataset (Dataset): 训练数据集（用于获取类别信息）
            label_mapping (dict): 标签映射字典
            num_classes (int): 类别数量
        
        Returns:
            float: 测试准确率
        """
        self.model.eval()
        running_corrects = 0
        total_samples = 0
        class_correct = [0] * num_classes
        class_total = [0] * num_classes
        
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(self.device), target.to(self.device)
                outputs = self.model(data)
                _, predicted = torch.max(outputs, 1)
                
                total_samples += target.size(0)
                running_corrects += (predicted == target).sum().item()
                
                # 计算每个类别的准确率
                c = (predicted == target).squeeze()
                for i in range(target.size(0)):
                    label = target[i]
                    if label < num_classes:
                        class_correct[label] += c[i].item()
                        class_total[label] += 1
        
        test_acc = 100. * running_corrects / total_samples
        print(f'\n测试集总体准确率: {test_acc:.2f}% ({running_corrects}/{total_samples})')
        
        # 打印每个类别的准确率
        if train_dataset and label_mapping:
            print('\n各类别测试准确率:')
            for i in range(min(num_classes, len(class_total))):
                if class_total[i] > 0:
                    class_acc = 100 * class_correct[i] / class_total[i]
                    # 获取类别名称
                    if i < len(train_dataset.classes):
                        class_id = train_dataset.classes[i]
                        class_name = label_mapping.get(class_id, 'unknown')
                    else:
                        class_name = 'unknown'
                    print(f'  {i:2d} ({class_name:15s}): {class_acc:6.2f}% '
                          f'({int(class_correct[i])}/{int(class_total[i])})')
        
        return test_acc
    
    def train(self, train_loader, val_loader, num_epochs, save_best=True):
        """
        完整训练过程
        
        Args:
            train_loader (DataLoader): 训练数据加载器
            val_loader (DataLoader): 验证数据加载器
            num_epochs (int): 训练轮数
            save_best (bool): 是否保存最佳模型
        
        Returns:
            dict: 训练历史
        """
        print(f'\n开始训练 {num_epochs} 个epoch...')
        start_time = time.time()
        
        best_val_acc = 0.0
        train_losses = []
        train_accs = []
        val_losses = []
        val_accs = []
        
        for epoch in range(num_epochs):
            print(f'\nEpoch {epoch+1}/{num_epochs}')
            print('-' * 50)
            
            # 训练
            train_loss, train_acc = self.train_epoch(train_loader, epoch)
            
            # 验证
            val_loss, val_acc = self.validate(val_loader)
            
            # 更新学习率
            current_lr = self.optimizer.param_groups[0]['lr']
            if self.scheduler:
                self.scheduler.step()
            
            # 记录结果
            train_losses.append(train_loss)
            train_accs.append(train_acc)
            val_losses.append(val_loss)
            val_accs.append(val_acc)
            
            # 记录到TensorBoard
            if self.tensorboard_logger:
                self.tensorboard_logger.log_training_epoch(
                    epoch, train_loss, train_acc, val_loss, val_acc, current_lr)
            
            print(f'训练 - 损失: {train_loss:.4f}, 准确率: {train_acc:.2f}%')
            print(f'验证 - 损失: {val_loss:.4f}, 准确率: {val_acc:.2f}%')
            print(f'当前学习率: {current_lr:.6f}')
            
            # 保存最佳模型
            if save_best and val_acc > best_val_acc:
                best_val_acc = val_acc
                if self.model_save_path:
                    save_model(self.model, self.model_save_path, epoch, val_acc, is_best=True)
                    print(f'保存最佳模型 (验证准确率: {best_val_acc:.2f}%)')
                
                # 记录最佳验证准确率到TensorBoard
                if self.tensorboard_logger:
                    self.tensorboard_logger.log_best_accuracy(best_val_acc, epoch)
        
        training_time = time.time() - start_time
        print(f'\n训练完成! 总用时: {training_time:.2f} 秒')
        print(f'最佳验证准确率: {best_val_acc:.2f}%')
        
        # 记录训练总结到TensorBoard
        if self.tensorboard_logger:
            batch_size = next(iter(train_loader))[0].size(0)
            initial_lr = self.optimizer.param_groups[0]['lr']
            self.tensorboard_logger.log_training_summary(
                training_time, best_val_acc, num_epochs, batch_size, initial_lr)
        
        return {
            'train_losses': train_losses,
            'train_accs': train_accs,
            'val_losses': val_losses,
            'val_accs': val_accs,
            'best_val_acc': best_val_acc,
            'training_time': training_time
        }
    
    def set_optimizer(self, optimizer):
        """设置优化器"""
        self.optimizer = optimizer
    
    def set_scheduler(self, scheduler):
        """设置学习率调度器"""
        self.scheduler = scheduler
    
    def set_criterion(self, criterion):
        """设置损失函数"""
        self.criterion = criterion 